package org.huangrui.spark.scala.core.acc

import org.apache.spark.util.AccumulatorV2
import org.apache.spark.{SparkConf, SparkContext}

import scala.collection.mutable

/**
 * @Author hr
 * @Create 2024-10-19 23:20 
 */
object Spark04_Acc_WordCount {
  def main(args: Array[String]): Unit = {
    val conf = new SparkConf().setMaster("local[*]").setAppName("spark")
    val sc = new SparkContext(conf)
    val rdd = sc.makeRDD(List("hello", "spark", "hello"))
    // 累加器 : WordCount
    // 创建累加器对象
    val wcAcc: MyAccumulator = new MyAccumulator
    // 向Spark进行注册
    sc.register(wcAcc)
    rdd.foreach(
      // 数据的累加（使用累加器）
      word => wcAcc.add(word)
    )
    // 获取累加器累加的结果
    println(wcAcc.value)

    sc.stop()
  }

  /**
   * 自定义数据累加器：WordCount
   * 1. 继承AccumulatorV2, 定义泛型
   * IN : 累加器输入的数据类型 String
   * OUT : 累加器返回的数据类型 mutable.Map[String, Int]
   * 2. 重写方法（6）
   */
  class MyAccumulator extends AccumulatorV2[String, mutable.Map[String, Int]] {
    private var wcMap: mutable.Map[String, Int] = mutable.Map[String, Int]()

    // 判断是否初始状态
    override def isZero: Boolean = wcMap.isEmpty

    override def copy(): AccumulatorV2[String, mutable.Map[String, Int]] = new MyAccumulator()

    override def reset(): Unit = wcMap.clear()

    // 获取累加器需要计算的值
    override def add(v: String): Unit = {
      wcMap.update(v, wcMap.getOrElse(v, 0) + 1)
    }

    //  Driver合并多个累加器
    override def merge(other: AccumulatorV2[String, mutable.Map[String, Int]]): Unit = {
      val map1 = wcMap
      val map2 = other.value
      map2.foreach({
        case (k, v) => map1.update(k, map1.getOrElse(k, 0) + v)
      })
    }
    // 累加器结果返回值
    override def value: mutable.Map[String, Int] = wcMap
  }
}
